package com.ml4ai.backend.utils;

import com.ml4ai.backend.dto.UserDTO;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;

import javax.websocket.Session;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Created by uesr on 2018/9/16.
 */
@Slf4j
public class WsHost {

    private static final Map<String, SessionStore> pool = new ConcurrentHashMap<>();

    @SneakyThrows
    public static boolean create(Session session) {
        pool.remove(session.getId());
        pool.put(session.getId(), new SessionStore(session, null));
        return true;
    }

    public static boolean assignUser(Session session, UserDTO user) {
        pool.get(session.getId()).setUser(user);
        return true;
    }

    @Data
    @AllArgsConstructor
    public static class SessionStore {
        private Session session;
        private UserDTO user;
    }

    public static void sendOnlineUser(Object message) {
        for (SessionStore store : pool.values()) {
            Session session = store.getSession();
            try {
                session.getAsyncRemote().sendObject(message);
            } catch (Exception e) {
                try {
                    pool.remove(session.getId());
                } catch (Exception ex) {
                    log.debug("{}", ex);
                } finally {
                    CloseUtils.close(session);
                }
            }
        }
    }

    public static void send2Users(List<Long> userIds, Object message) {
        if (userIds != null) {
            for (SessionStore store : pool.values()) {
                if (store.getUser() != null && userIds.contains(store.getUser().getId())) {
                    Session session = store.getSession();
                    try {
                        session.getAsyncRemote().sendObject(message);
                    } catch (Exception e) {
                        try {
                            pool.remove(session.getId());
                        } catch (Exception ex) {
                            log.debug("{}", ex);
                        } finally {
                            CloseUtils.close(session);
                        }
                    }
                }
            }
        }
    }
}
